#!/usr/bin/env python

import rospy
import numpy as np
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Odometry
from scipy.spatial.transform import Rotation as R

# 全局变量
all_error = 0.0
curr_error = 0.0

true_pose = np.zeros(3)
ltrue_pose = np.zeros(3)

odom_pose = np.zeros(3)
lodom_pose = np.zeros(3)

true_R = np.eye(3)
ltrue_R = np.eye(3)

odom_R = np.eye(3)
lodom_R = np.eye(3)

flag = False
time1 = None
time2 = None

def dongpucallback(pose):
    global true_pose, ltrue_pose, true_R, ltrue_R, time1, flag

    # 处理 true pose 数据
    rospy.logwarn("receive true pose")
    x = pose.pose.position.x
    y = pose.pose.position.y
    z = pose.pose.position.z

    qw = pose.pose.orientation.w
    qx = pose.pose.orientation.x
    qy = pose.pose.orientation.y
    qz = pose.pose.orientation.z

    ltrue_pose = true_pose
    ltrue_R = true_R

    true_pose = np.array([x, y, z])
    q = R.from_quat([qx, qy, qz, qw])
    true_R = q.as_matrix()

    time1 = rospy.Time.now()
    flag = True

def odomcallback(pose):
    global odom_pose, lodom_pose, odom_R, lodom_R, time2, flag, all_error, curr_error

    # 处理 odom pose 数据
    rospy.logwarn("receive odom pose")
    x = pose.pose.pose.position.x
    y = pose.pose.pose.position.y
    z = pose.pose.pose.position.z

    qw = pose.pose.pose.orientation.w
    qx = pose.pose.pose.orientation.x
    qy = pose.pose.pose.orientation.y
    qz = pose.pose.pose.orientation.z

    odom_pose = np.array([x, y, z])
    q = R.from_quat([qx, qy, qz, qw])
    odom_R = q.as_matrix()

    time2 = rospy.Time.now()
    if flag and (time2 >= time1):
        # 计算误差
        # 1. 平移误差
        t = (time1 - time2).to_sec()
        delt = np.linalg.norm(true_pose - ltrue_pose)

        vx = pose.twist.twist.linear.x
        vy = pose.twist.twist.linear.y
        vz = pose.twist.twist.linear.z

        wx = pose.twist.twist.angular.x
        wy = pose.twist.twist.angular.y
        wz = pose.twist.twist.angular.z

        odom_pose = odom_pose + t * np.array([vx, vy, vz])
        odelt = np.linalg.norm(odom_pose - lodom_pose)

        curr_error = (odelt - delt) ** 2
        rospy.logwarn("The translate error is %f", curr_error)

        # 2. 旋转误差
        delR = np.linalg.inv(ltrue_R) @ true_R
        delta_q = R.from_quat([wx * 0.5 * t, wy * 0.5 * t, wz * 0.5 * t, 1.0])
        q = q * delta_q
        odom_R = q.as_matrix()
        odelR = np.linalg.inv(lodom_R) @ odom_R
        error_R = np.linalg.inv(odelR) @ delR
        error_axis_angle = R.from_matrix(error_R).as_rotvec()
        rotation_error = np.linalg.norm(error_axis_angle)
        rospy.logwarn("The rotation error is %f", rotation_error)

        # all_error += curr_error + rotation_error
        all_error += curr_error
        rospy.logwarn("The total error is %f", all_error)

        lodom_pose = odom_pose
        lodom_R = odom_R
        flag = False
    else:
        rospy.logwarn("Waiting for true pose data...")

def main():
    rospy.init_node("trajectory_err", anonymous=True)
    rospy.logwarn("evo start")

    # 订阅话题
    rospy.Subscriber("/vrpn_client_node/uav1/pose", PoseStamped, dongpucallback)
    rospy.Subscriber("/baton/stereo3/odometry", Odometry, odomcallback)

    rospy.spin()

if __name__ == "__main__":
    try:
        main()
    except rospy.ROSInterruptException:
        pass